module Optim (getOptLevel,tests,runTests) where
import Common
import qualified Control.Monad.State as S
import qualified Data.IntMap as Map

import Test.HUnit


data Modification = Change Int | Reset Int deriving (Eq)
type SideEffect = (Offset, Map.IntMap Modification, [Int])

isStatic :: Operation -> Bool
isStatic op = case op of
   (_ :+= _)     -> True
   (_ :<- _)     -> True
   MoveOp _    -> True
   PutcOp (Const _) -> True
   PutcOp (Str _)  -> True
   _             -> False 

collectEffects :: [Operation] -> SideEffect -> SideEffect
collectEffects ops (off, assoc, outp) = case ops of
      [] -> (off, assoc,outp)
      (MoveOp m):xs -> collectEffects xs (off + m, assoc, outp) 
      (sof :<- v):xs -> collectEffects xs (off, Map.insert (off + sof) (Reset v) assoc, outp)
      (PutcOp (Const i)):xs -> collectEffects xs (off, assoc, outp ++ [i])
      (PutcOp (Str s)):xs -> collectEffects xs (off, assoc, outp ++ s)
      (iof :+= v):xs ->
          let noff = iof + off 
          in if Map.member noff assoc
              then collectEffects xs (off, Map.adjust (modMod v) noff assoc, outp)
              else collectEffects xs (off, Map.insert noff (Change v) assoc, outp)
      _  -> error "Somehow, a non-static got in collectEffects."

modMod :: Int -> Modification -> Modification 
modMod i (Change v) = Change (v + i)
modMod i (Reset v) = Reset (v + i)


frontEffects :: Transformation
frontEffects lst = effectsToOps (roff, (Map.map initFix assoc), outp) ++ findEffects rest
 where (front,rest) = span isStatic lst
       (roff, assoc, outp) = collectEffects front (0, Map.empty, [])
       initFix (Change v) = Reset v
       initFix x          = x

findEffects :: Transformation
findEffects [] = []
findEffects lst = 
   mapOB findEffects front ++ 
    effectsToOps (collectEffects valid (0, Map.empty, [])) 
    ++ findEffects rest
 where (front, more) = break isStatic lst 
       (valid,rest) = span isStatic more 

mapOB :: Transformation -> [Operation] -> [Operation]
mapOB _ []                 = []
mapOB f ((LoopOp t b):xs)  = LoopOp t (f b) : mapOB f xs
mapOB f (x:xs)             = x : mapOB f xs 


effectsToOps :: SideEffect -> [Operation]
effectsToOps (off, assoc, outp) = (toPutcOp outp) ++ map conv pairs ++ rest
 where pairs = Map.assocs assoc
       conv (offs, Change i) = offs :+= i
       conv (offs, Reset i)  = offs :<- i
       rest = if off == 0 then [] else [MoveOp off]
       toPutcOp [] = []
       toPutcOp [x] = [PutcOp (Const x)]
       toPutcOp lst = [PutcOp (Str lst)]
   --    isLoop = Map.lookup 0 assoc == Just (Change (-1))


type Optimization = ([Operation] -> S.State Bool [Operation])

evalOpt ofun ops = S.evalState (wrap (ofun ops)) False
 where wrap f = do res <- f
                   s <- S.get
                   return (res,s)

runOptims :: [Operation] -> [Optimization] -> [Operation]
runOptims tree opts = if changed then runOptims tr opts else tr
 where (tr,changed) = evalOpt bigFun tree
       bigFun :: Optimization
       bigFun = foldl ffun return opts 
       ffun a b = \x -> b x >>= a

runOptimizations :: Transformation
runOptimizations tree = runOptims tree optims
 where optims :: [Optimization]
       optims =  [ groupSim, cleanBlocks,
                    addsToSets, optAdds,
                    optimOne, removeDups, 
                    doAdds, reduceOp1,
                    optimOne, optAdds, 
                    removeDups, addsToSets
                 ]

removeDups ((MoveOp 0):xs) = markChange >> removeDups xs 
removeDups ((_ :+= 0):xs)  = markChange >> removeDups xs
removeDups (a:b:xs) = case (a,b) of
  (o1 :<- _, o2 :<- _) | o1 == o2 -> markChange >> rest 
  (o1 :+= _, o2 :<- _) | o1 == o2 -> markChange >> rest 
  (o1 :+= v1, o2 :+= v2) | o1 == o2 -> nextWith (o1 :+= (v1+v2))
  (o1 :<- v1, o2 :+= v2) | o1 == o2 -> nextWith (o1 :<- (v1+v2))
  (OpLocAddT o1 _ _, o2 :<- _) | o1 == o2 -> markChange >> rest
  (MoveOp m, MoveOp n) -> nextWith (MoveOp (m+n))
  (o1 :+= v1, MoveOp o2) | o1 == o2 -> markChange >> (retCons (MoveOp o1) $ (removeDups ((0 :+= v1):xs)))
  _                    -> continue
 where continue = (removeDups `onLoopOp` a) >>= (`retCons` rest)
       rest     = removeDups (b:xs)
       nextWith v = markChange >> removeDups (v:xs)
removeDups lst = handleRest removeDups lst


onLoopOp f (LoopOp t b) = f b >>= \bp -> return (LoopOp t bp)
onLoopOp _ x            = return x

retCons a b =  do
  rest <- b
  return (a:rest)

retCat a b =  do
  rest <- b
  return (a ++ rest)

markChange = S.put True

groupSim :: Optimization
groupSim ((MoveOp o1):(OpLocAddT o2 o3 m):xs) = markChange >> (retCons (OpLocAddT (o2+o1) (o3+o1) m) $ groupSim ((MoveOp o1):xs))
groupSim ((MoveOp o1):(OpLocSetT o2 o3 m):xs) = markChange >> (retCons (OpLocSetT (o2+o1) (o3+o1) m) $ groupSim ((MoveOp o1):xs))
--groupSim ((MoveOp o1):(o2 :<- v):xs) = markChange >> (retCons ((o2+o1) :<- v) $ groupSim ((MoveOp o1):xs))
groupSim (a:b:c:xs) = case (a,b,c) of
 (MoveOp o1, o3 :+= m, MoveOp o2) -> 
      markChange >> (retCons ((o1+o3) :+= m) $ groupSim (MoveOp (o1+o2):xs))
 (MoveOp o1, o3 :<- m, MoveOp o2) -> 
      markChange >> (retCons ((o1+o3) :<- m) $ groupSim (MoveOp (o1+o2):xs))
 _ -> handleRest groupSim (a:b:c:xs)
groupSim lst = handleRest groupSim lst


doAdds ((LoopOp t [0 :+= (-1), p :+= o]):xs) =
 if p /= 0  -- (current + p) += current * o 
  then markChange >> (retCat [OpLocAddT p 0 o, 0 :<- 0] $ doAdds xs)
  else markChange >> doAdds ((LoopOp t [0 :+= (o-1)]):xs)
doAdds ((LoopOp t [p :+= o, 0 :+= (-1)]):xs) =
 if p /= 0  
  then markChange >> (retCat [OpLocAddT p 0 o, 0 :<- 0] $ doAdds xs)
  else markChange >> doAdds ((LoopOp t [0 :+= (o-1)]):xs) 
doAdds ((o1 :<- 0):(MoveOp o2):(OpLocAddT o3 0 m):xs) 
 | o1 - o2 == o3 = markChange >> doAdds ((OpLocSetT (o2+o3) o2 m):(MoveOp o2):xs)
doAdds lst = handleRest doAdds lst
  

reduceOp1 ((o1 :<- v1):(o2 :+= v2):xs) 
   | o1 == o2 = markChange >> reduceOp1 ((o1 :<- (v1+v2)):xs) 
reduceOp1 ((0 :<- v):(PutcOp (Ptr 0)):xs) = markChange >> (retCons (PutcOp (Const v)) $ reduceOp1 ((0 :<- v):xs))
reduceOp1 ((o1 :<- v):(PutcOp (Ptr o2)):xs) = 
  if o1 == o2 
   then if v < 0 then error (show (o1 :<- v))
                 else markChange >> (retCons (PutcOp (Const v)) $ reduceOp1 ((o1 :<- v):xs))
   else markChange >> (retCons (PutcOp (Ptr o2)) $ reduceOp1 ((o1 :<- v):xs))
reduceOp1 ((0 :<- v):(PutcOp (Const i)):xs) = markChange >> (retCons (PutcOp (Const i)) $ reduceOp1 ((0 :<- v):xs))

reduceOp1 ((v@(MoveOp _)):(pc@(PutcOp (Const _))):xs) = markChange >> reduceOp1 (pc:v:xs)
reduceOp1 ((v@(_ :<- _)):(pc@(PutcOp (Const _))):xs) = markChange >> reduceOp1 (pc:v:xs)
reduceOp1 ((v@(_ :+= _)):(pc@(PutcOp (Const _))):xs) = markChange >> reduceOp1 (pc:v:xs)

reduceOp1 ((PutcOp (Const i1)):(PutcOp (Const i2)):xs) = markChange >> reduceOp1 ((PutcOp (Str [i1,i2])):xs)
reduceOp1 ((PutcOp (Str lst)):(PutcOp (Const i2)):xs) = markChange >> reduceOp1 ((PutcOp (Str (lst ++ [i2]))):xs)
reduceOp1 ((PutcOp (Const i1)):(PutcOp (Str lst)):xs) = markChange >> reduceOp1 ((PutcOp (Str (i1:lst))):xs)

reduceOp1 ((0 :<- v):(OpLocAddT o 0 m):xs) = markChange >> (retCons (o :+= (v*m)) $ reduceOp1 ((0 :<- v):xs))
reduceOp1 ((0 :<- v):(OpLocSetT o 0 m):xs) = markChange >> (retCons (o :<- (v*m)) $ reduceOp1 ((0 :<- v):xs))
reduceOp1 ((o1 :<- v):(OpLocAddT o2 o3 m):xs) 
 | o1 == o3 = markChange >> (retCons (o2 :+= (v*m)) $ reduceOp1 ((o1 :<- v):xs))
 | o1 == o2 && v == 0 = markChange >> reduceOp1 ((OpLocSetT o2 o3 m):xs)
reduceOp1 ((o1 :<- v):(OpLocSetT o2 o3 m):xs)
 | o1 == o3 = markChange >> (retCons (o2 :<- (v*m)) $ reduceOp1 ((o1 :<- v):xs))
reduceOp1 ((MoveOp o1):(PutcOp (Ptr 0)):(MoveOp o2):xs) = markChange >> (retCons (PutcOp (Ptr o1)) $ reduceOp1 ((MoveOp (o1+o2)):xs))
reduceOp1 ((MoveOp o1):(PutcOp (Ptr o2)):xs) = markChange >> (retCons (PutcOp (Ptr (o1+o2))) $ reduceOp1 ((MoveOp o1):xs))
reduceOp1 ((off1 :<- v1):(MoveOp off2):xs) 
  | off1 == off2 = markChange >> reduceOp1 ((MoveOp off2):( 0 :<- v1 ):xs)
reduceOp1 lst = handleRest reduceOp1 lst


cleanBlocks ((0 :<- v):(LoopOp While b):xs) = 
 if v /= 0 
  then markChange >> (retCons (0 :<- v) $ cleanBlocks ((LoopOp DoWhile b):xs))
  else markChange >> cleanBlocks ((0 :<- 0):xs) 
cleanBlocks ((MoveOp o1):(LoopOp While [MoveOp o2]):xs) = 
 if o1 == o2 
   then markChange >> cleanBlocks ((LoopOp DoWhile [MoveOp o2]):xs) -- FIXME
   --then retCons (MoveOp o1) $ cleanBlocks ((LoopOp While [MoveOp o2]):xs)
   else retCons (MoveOp o1) $ cleanBlocks ((LoopOp While [MoveOp o2]):xs)
cleanBlocks (b@(LoopOp _ _):(0 :+= m):xs) = markChange >> cleanBlocks (b:(0 :<- m):xs) 
cleanBlocks (b@(LoopOp _ _):(OpLocAddT 0 off m):xs) = markChange >> cleanBlocks (b:(OpLocSetT 0 off m):xs) 
cleanBlocks (b@(LoopOp t _):(0 :<- 0):xs) | t /= If  = markChange >> cleanBlocks (b:xs)
cleanBlocks ((LoopOp t x):xs) = do
  bcleaned <- cleanBlocks x
  if t == While && not (null bcleaned) && last bcleaned == (0 :<- 0)
    then markChange >> (retCons (LoopOp If bcleaned) $ cleanBlocks xs)
    else retCons (LoopOp t bcleaned) $ cleanBlocks xs
cleanBlocks (x:xs) = retCons x $ cleanBlocks xs 
cleanBlocks [] = return []

optimOne (MoveOp off:OpLocAddT off2 so m:xs) | off + off2 == 0 
   = markChange >> optimOne (OpLocAddT (off + off2) (so+off) m:MoveOp off:xs)
optimOne (a:b:c:d:xs) = case (a,b,c,d) of
 (MoveOp off1, (OpLocAddT off 0 m), (0 :<- 0), (MoveOp off2)) ->
   if (off1+off) == 0 
      then markChange >> optimOne ((OpLocAddT 0 off1 m):(off1 :<- 0):(MoveOp (off1+off2)):xs) 
      else retCat [a,b,c] $ optimOne (d:xs) 
 (_,_,_,_) -> handleRest optimOne (a:b:c:d:xs)
optimOne lst = handleRest optimOne lst


isAssign (_ :<- _) = True
isAssign _         = False

-- TODO: Effect analysis in the loop to check for simple decrement
optAdds :: Optimization
optAdds ((LoopOp While ((0 :+= (-1)):bx)):xs)
  | all liftable bx =
       let res = ifify $ findEffects (map fixup bx)
       in markChange >> retCat (res ++ [0 :<- 0]) (optAdds xs) 
 where liftable (off :+= _) = off /= 0
       liftable (off :<- _) = off /= 0
       liftable (PutcOp (Const _)) = True 
       liftable _           = False
       ifify lst = if hasAssign lst then [LoopOp If lst] else lst
       hasAssign lst = any isAssign lst
       fixup (off :+= m) = OpLocAddT off 0 m
       fixup (off :<- x) = off :<- x
       fixup _           = error "Bad param in fixup"
optAdds lst = handleRest optAdds lst


handleRest f lst = case lst of
  (x:xs)             -> (f `onLoopOp` x) >>= (`retCons` (f xs))
  []                 -> return []

addsToSets ((idx :<- 0):(OpLocAddT idx2 off m) : xs)
  | idx == idx2 = markChange >> retCons (OpLocSetT idx off m) (addsToSets xs)
  | otherwise = retCat [idx :<- 0, OpLocAddT idx2 off m] (addsToSets xs)
addsToSets lst = handleRest addsToSets lst

run_addsToSets l = fst (evalOpt addsToSets l)

clean = filter (not . isNop)
 where isNop op = case op of 
                     MoveOp 0  -> True
                     (_ :+= 0) -> True
                     _         -> False

trim lst = reverse $ ((filter isIo) back) ++ rest
 where (back,rest) = span safe (reverse lst)
       noDep (PutcOp (Const _)) = True
       noDep (PutcOp (Str _)) = True
       noDep GetcOp = True
       noDep _      = False
       safe x = noDep x || isStatic x
       isIo (PutcOp _) = True
       isIo GetcOp = True
       isIo _          = False


fixFirst [] = []
fixFirst ((_ :<- 0):xs) = fixFirst xs
fixFirst ((n :+= v):xs) = fixFirst ((n :<- v):xs)
fixFirst ((LoopOp While _):xs) = fixFirst xs
fixFirst ((LoopOp If _):xs) = fixFirst xs
fixFirst ((MoveOp x):xs) = (MoveOp x):(0 :<- 0):xs
fixFirst v = v

getOptLevel :: Int -> Transformation
getOptLevel n = trim . clean . fun
 where fun = case n of
        0 -> fixFirst 
        1 -> run_addsToSets . getOptLevel 0
        2 -> frontEffects . getOptLevel 1
        3 -> run_addsToSets . getOptLevel 2
        4 -> runOptimizations . getOptLevel 3
        5 -> frontEffects . getOptLevel 4 . frontEffects .  getOptLevel 4
        6 -> fixFirst . runOptimizations . getOptLevel 5
        7 -> fixFirst . frontEffects . runOptimizations . frontEffects . runOptimizations . frontEffects . getOptLevel 6
        _ -> error "Invalid Optimization level"


-- TESTS -- 
gotest = runTestTT tests
runTests = gotest

tests = TestList [
    testRemoveDups
   ,testFindEffects
   ,testAddsToSets
   ,testOptAdds
   ,testOptimOne
   ,testWillOpt
 ]

testFindEffects = "findEffects" ~: TestList [
     [MoveOp 0] `becomes` []
     ,[0 :<- 2, 0 :<- 4] `becomes` [0 :<- 4]
     ,[0 :+= 2, 0 :+= 4] `becomes` [0 :+= 6]
     ,[0 :+= 2, 1 :<- 4, 0 :<- 5] `becomes` [0 :<- 5, 1 :<- 4]
     ,[MoveOp 1,0 :<- 9,MoveOp 1, (-1) :<- 4, MoveOp (-2), 1 :+= (-4)] `becomes` [1 :<- 0]
   ]
 where becomes a res = res ~=? findEffects a

testOptimOne = "optimOne" ~: TestList [
    [MoveOp 1, OpLocAddT (-1) 0 10] `becomes` [OpLocAddT 0 1 10, MoveOp 1]
  ]
 where becomes a res = (res,True) ~=? evalOpt optimOne a

testOptAdds = "optAdds" ~: TestList [
     "nop while" ~: whileLoop [] `becomes` [0 :<- 0]
     ,whileLoop [1 :<- 9] `becomes` [ifOp [1 :<- 9], 0 :<- 0]
     ,whileLoop [1 :<- 9, 1 :<- 12] `becomes` [ifOp [1 :<- 12], 0 :<- 0]
     ,whileLoop [2 :<- 3, 2 :+= 4] `becomes` [ifOp [2 :<- 3, OpLocAddT 2 0 4],  0 :<- 0]
     ,[0 :<- 0, LoopOp While [0 :+= (-1), 2 :<- 2]] `becomes` [0 :<- 0, ifOp [2 :<- 2], 0 :<- 0]
     ,no_change (whileLoop [1 :<- 4, 0 :<- 0]) 
   ]
 where becomes a res = (res,True) ~=? evalOpt optAdds a
       whileLoop ops = [LoopOp While (0 :+= (-1):ops)]
       ifOp ops = LoopOp If ops
       no_change x = (x,False) ~=? evalOpt optAdds x

testAddsToSets = "addsToSets" ~: TestList [
     [3 :<- 0, OpLocAddT 3 2 1] `becomes` [OpLocSetT 3 2 1]
     ,no_change [4 :<- 0, OpLocAddT 3 2 1]
   ]
 where becomes a res = (res,True) ~=? evalOpt addsToSets a
       no_change a = (a,False) ~=? evalOpt addsToSets a

testRemoveDups = "removeDups" ~: TestList [
     [MoveOp 0] `becomes` []
     ,[0 :<- 2, 0 :<- 4] `becomes` [0 :<- 4]
     ,[0 :+= 2, 0 :+= 4] `becomes` [0 :+= 6]
   ]
 where becomes a res = (res,True) ~=? evalOpt removeDups a

testWillOpt = "big opt" ~: TestList [
     [0 :<- 3, 0 :+= 40, PutcOp (Ptr 0)] `becomes` [PutcOp (Const 43)]
   ]
 where becomes a res = res ~=? getOptLevel 7 a

